package step

import (
	"bytes"
	"crypto/ecdsa"
	"errors"
	"fmt"
	"github.com/wanchain/go-mpc/common/hexutil"
	"math/big"
	"strconv"

	"github.com/wanchain/go-mpc/common"
	"github.com/wanchain/go-mpc/crypto"
	"github.com/wanchain/go-mpc/log"
	mpcprotocol "github.com/wanchain/go-mpc/storeman/storemanmpc/protocol"
)

type TxSign_CalSignStepUni struct {
	TXSign_Lagrange_StepUni
	signNum int
}

func CreateTxSign_CalSignStepUni(totalNum int, resultKey string, signNum int) *TxSign_CalSignStepUni {
	log.SyslogInfo("CreateTxSign_CalSignStepUni begin")

	signSeedKeys := mpcprotocol.GetPreSetKeyArr(mpcprotocol.MpcTxSignSeed, signNum)
	resultKeys := mpcprotocol.GetPreSetKeyArr(resultKey, signNum)
	mpc := &TxSign_CalSignStepUni{*CreateTXSign_Lagrange_StepUni(totalNum, signSeedKeys, resultKeys), signNum}
	return mpc
}

func (txStep *TxSign_CalSignStepUni) GetMsgGens() []MpcMessageGenerator {
	return txStep.TXSign_Lagrange_StepUni.GetMsgGens()
}

func (txStep *TxSign_CalSignStepUni) InitStep(result mpcprotocol.MpcResultInterface) error {
	log.SyslogDebug("TxSign_CalSignStepUni::InitStep begin", "ctx", txStep.GetMpcContextId())

	privateKey, err := result.GetValue(mpcprotocol.MpcPrivateShare)
	if err != nil {
		return err
	}

	for i := 0; i < txStep.signNum; i++ {
		ar, err := result.GetValue(mpcprotocol.MpcSignARResult + "_" + strconv.Itoa(i))
		if err != nil {
			return err
		}

		aPoint, err := result.GetValue(mpcprotocol.MpcSignAPoint + "_" + strconv.Itoa(i))
		if err != nil {
			return err
		}

		r, err := result.GetValue(mpcprotocol.MpcSignR + "_" + strconv.Itoa(i))
		if err != nil {
			return err
		}

		c, err := result.GetValue(mpcprotocol.MpcSignC + "_" + strconv.Itoa(i))
		if err != nil {
			return err
		}

		txHash := txStep.GetMpcContext().GetInput().HashData[i]

		arInv := ar[0]
		arInv.ModInverse(&arInv, crypto.Secp256k1_N)
		invRPoint := new(ecdsa.PublicKey)
		invRPoint.Curve = crypto.S256()
		invRPoint.X, invRPoint.Y = crypto.S256().ScalarMult(&aPoint[0], &aPoint[1], arInv.Bytes())
		if invRPoint.X == nil || invRPoint.Y == nil {
			log.SyslogErr("TxSign_CalSignStepUni::InitStep, invalid r point", "ctx", txStep.GetMpcContextId())
			return mpcprotocol.ErrPointZero
		}

		log.SyslogDebug(fmt.Sprintf("TxSign_CalSignStepUni::InitStep, calsign, x:%s, y:%s", invRPoint.X.String(), invRPoint.Y.String()), "ctx", txStep.GetMpcContextId())
		SignSeed := new(big.Int).Set(invRPoint.X)
		SignSeed.Mod(SignSeed, crypto.Secp256k1_N)
		var v int64
		if invRPoint.X.Cmp(SignSeed) == 0 {
			v = 0
		} else {
			v = 2
		}

		invRPoint.Y.Mod(invRPoint.Y, big.NewInt(2))
		if invRPoint.Y.Cmp(big.NewInt(0)) != 0 {
			v |= 1
		}

		log.SyslogDebug(fmt.Sprintf("TxSign_CalSignStepUni::InitStep, %s:%s, %s:%d",
			mpcprotocol.MpcTxSignResultR+"_"+strconv.Itoa(i), SignSeed.String(), mpcprotocol.MpcTxSignResultV+"_"+strconv.Itoa(i), v), "ctx", txStep.GetMpcContextId())
		result.SetValue(mpcprotocol.MpcTxSignResultR+"_"+strconv.Itoa(i), []big.Int{*SignSeed})
		result.SetValue(mpcprotocol.MpcTxSignResultV+"_"+strconv.Itoa(i), []big.Int{*big.NewInt(v)})
		SignSeed.Mul(SignSeed, &privateKey[0])
		SignSeed.Mod(SignSeed, crypto.Secp256k1_N)

		//hash := txHash[0]
		hashBig := big.NewInt(0).SetBytes(txHash)

		SignSeed.Add(SignSeed, hashBig)
		SignSeed.Mod(SignSeed, crypto.Secp256k1_N)
		SignSeed.Mul(SignSeed, &r[0])
		SignSeed.Mod(SignSeed, crypto.Secp256k1_N)
		SignSeed.Add(SignSeed, &c[0])
		SignSeed.Mod(SignSeed, crypto.Secp256k1_N)

		result.SetValue(mpcprotocol.MpcTxSignSeed+"_"+strconv.Itoa(i), []big.Int{*SignSeed})
		log.SyslogInfo("TxSign_CalSignStepUni::InitStep", "S seed", SignSeed.String(), "ctx", txStep.GetMpcContextId(), "hash", hexutil.Encode(txHash))
		log.SyslogDebug(fmt.Sprintf("TxSign_CalSignStepUni::InitStep, %s:%s", mpcprotocol.MpcTxSignSeed+"_"+strconv.Itoa(i), SignSeed.String()), "ctx", txStep.GetMpcContextId())
	}

	err = txStep.TXSign_Lagrange_StepUni.InitStep(result)
	if err != nil {
		log.SyslogErr(fmt.Sprintf("TxSign_CalSignStepUni::InitStep, initStep fail, err:%s", err.Error()), "ctx", txStep.GetMpcContextId())
		return err
	} else {
		log.SyslogDebug("TxSign_CalSignStepUni::InitStep succeed", "ctx", txStep.GetMpcContextId())
		return nil
	}
}

func (txStep *TxSign_CalSignStepUni) FinishStep(result mpcprotocol.MpcResultInterface, mpc mpcprotocol.StoremanManager) error {
	log.SyslogDebug("TxSign_CalSignStepUni::FinishStep begin", "ctx", txStep.GetMpcContextId())

	err := txStep.TXSign_Lagrange_StepUni.FinishStep(result, mpc)
	if err != nil {
		return err
	}

	err = txStep.VerifyRS(result, mpc)
	if err != nil {
		return err
	}

	txStep.TXSign_Lagrange_StepUni.SetSOKIndex()

	log.SyslogDebug("TxSign_CalSignStepUni::FinishStep succeed", "ctx", txStep.GetMpcContextId())
	return nil
}

func (txStep *TxSign_CalSignStepUni) VerifyRS(result mpcprotocol.MpcResultInterface, mpc mpcprotocol.StoremanManager) error {
	for i := 0; i < txStep.signNum; i++ {
		iStr := "_" + strconv.Itoa(i)

		R, err := result.GetValue(mpcprotocol.MpcTxSignResultR + iStr)
		if err != nil {
			return err
		}

		V, err := result.GetValue(mpcprotocol.MpcTxSignResultV + iStr)
		if err != nil {
			return err
		}

		S, err := result.GetValue(mpcprotocol.MpcTxSignResult + iStr)
		if err != nil {
			return err
		}

		hashBytes := txStep.MpcContext.GetInput().HashData[i]

		pk, err := RecoverPublicKey(hashBytes[:], &R[0], &S[0], &V[0], txStep.GetMpcContextId())
		if err != nil {
			log.SyslogErr("TxSign_CalSignStepUni::VerifyRS", "err", err.Error(), "ctx", txStep.GetMpcContextId())
			return err
		}
		if !bytes.Equal(pk[:], txStep.MpcContext.GetInput().PKBytes) {
			log.SyslogErr("TxSign_CalSignStepUni::VerifyRS", "err", mpcprotocol.ErrRecoverPkNotMatchPre.Error(), "ctx", txStep.GetMpcContextId())
			return mpcprotocol.ErrRecoverPkNotMatchPre
		}

		var buf bytes.Buffer
		buf.Write(common.LeftPadBytes(R[0].Bytes(), 32))
		buf.Write(common.LeftPadBytes(S[0].Bytes(), 32))
		buf.Write(common.LeftPadBytes(V[0].Bytes(), 1))

		key := mpcprotocol.MpcContextResult + strconv.Itoa(i)
		result.SetByteValue(key, buf.Bytes())
	}
	log.SyslogInfo("TxSign_CalSignStepUni", "@@@@@@@VerifyRS success", txStep.MpcContext.GetInput().String(), "ctx", txStep.GetMpcContextId())
	return nil
}

func RecoverPublicKey(sighash []byte, R, S, Vb *big.Int, ctxId uint64) ([]byte, error) {
	log.SyslogInfo("TxSign_CalSignStepUni::RecoverPublicKey", "Hash", common.ToHex(sighash[:]), "R", common.ToHex(R.Bytes()), "S", common.ToHex(S.Bytes()), "V", common.ToHex(Vb.Bytes()), "ctx", ctxId)

	if Vb.BitLen() > 8 {
		log.SyslogErr("TxSign_CalSignStepUni::RecoverPublicKey, invalid sign", "ctx", ctxId)
		return nil, errors.New("invalid sign")
	}

	// encode the snature in uncompressed format
	r, s := R.Bytes(), S.Bytes()
	sig := make([]byte, 65)
	copy(sig[32-len(r):32], r)
	copy(sig[64-len(s):64], s)

	vBytes := Vb.Bytes()
	if len(vBytes) > 0 {
		sig[64] = vBytes[0]
	}

	pubKey, err := crypto.SigToPub(sighash[:], sig)
	if err != nil {
		log.SyslogErr("TxSign_CalSignStepUni::RecoverPublicKey, fail", "err", err.Error(), "ctx", ctxId)
		return nil, err
	}

	return crypto.FromECDSAPub(pubKey)[1:], nil
}
